import os
import base64
import io
import numpy as np
from flask import Flask, request, jsonify
from flask_restful import Api, Resource
from PIL import Image
import cv2
import face_recognition

from src.algorithms.face_recognition import FaceRecognition
from src.utils.face_utils import create_faces_directory, put_chinese_text

# 初始化Flask应用
app = Flask(__name__)
api = Api(app)

# 创建人脸识别实例
face_recognizer = FaceRecognition()

# 创建faces目录
faces_dir = create_faces_directory()
print(f"人脸图像目录: {faces_dir}")

# 加载已知人脸
face_recognizer.load_known_faces(faces_dir)

class HealthCheck(Resource):
    """健康检查接口"""
    
    def get(self):
        return {'status': 'ok', 'message': 'Face recognition service is running'}

class RecognizeImage(Resource):
    """图片人脸识别接口"""
    
    def post(self):
        try:
            # 检查是否有上传文件
            if 'image' not in request.files:
                return {'error': 'No image file provided'}, 400
            
            file = request.files['image']
            if file.filename == '':
                return {'error': 'No image file selected'}, 400
            
            # 读取图像文件
            image_bytes = file.read()
            np_arr = np.frombuffer(image_bytes, np.uint8)
            image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
            
            if image is None:
                return {'error': 'Invalid image file'}, 400
            
            # 转换BGR到RGB
            rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            
            # 找到图片中所有人脸的位置和编码
            face_locations = face_recognition.face_locations(rgb_image)
            face_encodings = face_recognition.face_encodings(rgb_image, face_locations)
            
            face_names = []
            for face_encoding in face_encodings:
                matches = face_recognition.compare_faces(
                    face_recognizer.known_face_encodings, face_encoding)
                name = "Unknown"
                
                if len(face_recognizer.known_face_encodings) > 0:
                    face_distances = face_recognition.face_distance(
                        face_recognizer.known_face_encodings, face_encoding)
                    best_match_index = np.argmin(face_distances)
                    
                    if matches[best_match_index]:
                        name = face_recognizer.known_face_names[best_match_index]
                
                face_names.append(name)
            
            # 返回结果
            result = []
            for (top, right, bottom, left), name in zip(face_locations, face_names):
                result.append({
                    'name': name,
                    'location': {
                        'top': top,
                        'right': right,
                        'bottom': bottom,
                        'left': left
                    }
                })
            
            return {
                'success': True,
                'faces_detected': len(result),
                'faces': result
            }
        
        except Exception as e:
            return {'error': f'Failed to process image: {str(e)}'}, 500

class RecognizeBase64Image(Resource):
    """Base64图片人脸识别接口"""
    
    def post(self):
        try:
            # 获取JSON数据
            data = request.get_json()
            if not data or 'image' not in data:
                return {'error': 'No image data provided'}, 400
            
            # 解码Base64图像
            image_data = data['image']
            if image_data.startswith('data:image'):
                # 移除数据URI前缀
                image_data = image_data.split(',')[1]
            
            image_bytes = base64.b64decode(image_data)
            np_arr = np.frombuffer(image_bytes, np.uint8)
            image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
            
            if image is None:
                return {'error': 'Invalid image data'}, 400
            
            # 转换BGR到RGB
            rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            
            # 找到图片中所有人脸的位置和编码
            face_locations = face_recognition.face_locations(rgb_image)
            face_encodings = face_recognition.face_encodings(rgb_image, face_locations)
            
            face_names = []
            for face_encoding in face_encodings:
                matches = face_recognition.compare_faces(
                    face_recognizer.known_face_encodings, face_encoding)
                name = "Unknown"
                
                if len(face_recognizer.known_face_encodings) > 0:
                    face_distances = face_recognition.face_distance(
                        face_recognizer.known_face_encodings, face_encoding)
                    best_match_index = np.argmin(face_distances)
                    
                    if matches[best_match_index]:
                        name = face_recognizer.known_face_names[best_match_index]
                
                face_names.append(name)
            
            # 返回结果
            result = []
            for (top, right, bottom, left), name in zip(face_locations, face_names):
                result.append({
                    'name': name,
                    'location': {
                        'top': top,
                        'right': right,
                        'bottom': bottom,
                        'left': left
                    }
                })
            
            return {
                'success': True,
                'faces_detected': len(result),
                'faces': result
            }
        
        except Exception as e:
            return {'error': f'Failed to process image: {str(e)}'}, 500

class AddFace(Resource):
    """添加新人脸接口"""
    
    def post(self):
        try:
            # 获取JSON数据
            data = request.get_json()
            if not data or 'name' not in data or 'image' not in data:
                return {'error': 'Missing name or image data'}, 400
            
            name = data['name']
            image_data = data['image']
            
            # 解码Base64图像
            if image_data.startswith('data:image'):
                # 移除数据URI前缀
                image_data = image_data.split(',')[1]
            
            image_bytes = base64.b64decode(image_data)
            np_arr = np.frombuffer(image_bytes, np.uint8)
            image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
            
            if image is None:
                return {'error': 'Invalid image data'}, 400
            
            # 转换BGR到RGB
            rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            face_encodings = face_recognition.face_encodings(rgb_image)
            
            if len(face_encodings) == 0:
                return {'error': 'No face detected in the image'}, 400
            
            # 添加人脸编码和名称
            face_recognizer.known_face_encodings.append(face_encodings[0])
            face_recognizer.known_face_names.append(name)
            
            # 保存图像文件到faces目录
            filename = f"{name}.jpg"
            target_path = os.path.join(faces_dir, filename)
            # 处理中文文件名保存问题
            is_success, im_buf_arr = cv2.imencode(".jpg", image)
            if is_success:
                im_buf_arr.tofile(target_path)
            else:
                return {'error': 'Failed to encode image'}, 500
            
            return {
                'success': True,
                'message': f'Successfully added face: {name}',
                'name': name
            }
        
        except Exception as e:
            return {'error': f'Failed to add face: {str(e)}'}, 500

class ListFaces(Resource):
    """列出所有已知人脸接口"""
    
    def get(self):
        try:
            faces = []
            for filename in os.listdir(faces_dir):
                if filename.lower().endswith(('.jpg', '.png', '.jpeg')):
                    name = os.path.splitext(filename)[0]
                    faces.append({
                        'name': name,
                        'image': filename
                    })
            
            return {
                'success': True,
                'faces_count': len(faces),
                'faces': faces
            }
        
        except Exception as e:
            return {'error': f'Failed to list faces: {str(e)}'}, 500

# 添加路由
api.add_resource(HealthCheck, '/api/health')
api.add_resource(RecognizeImage, '/api/recognize/image')
api.add_resource(RecognizeBase64Image, '/api/recognize/base64')
api.add_resource(AddFace, '/api/faces')
api.add_resource(ListFaces, '/api/faces')

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=True)